from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory
from rest_framework.response import Response
from rest_framework.views import APIView

class AskAPIView(APIView):
    def __init__(self):
        super().__init__()
        self.memory = ChatMessageHistory()

    def post(self, request):
        # 初始化模型和提示模板
        llm = Tongyi()

        # 验证请求数据
        if 'input' not in request.data:
            return Response({'msg': 'error', 'data': '缺少 input 参数'}, status=400)

        input = request.data['input']

        # 构建提示模板
        prompt = ChatPromptTemplate.from_messages([
            ('system', "你是一名资深医生,可以回答患者的问题"),
            MessagesPlaceholder(variable_name="history"),
            ("human", "{input}")
        ])
        chain = prompt | llm

        def get_memory():
            return self.memory

        runnable_history = RunnableWithMessageHistory(
            chain,
            get_memory,
            input_messages_key='input',
            history_messages_key='history',
        )

        try:
            # 使用历史消息构建上下文并生成回复
            res = runnable_history.invoke({'input': input})

            # 将对话的消息添加到历史记录中
            self.memory.add_user_message(input)
            self.memory.add_ai_message(res)

            print(res)
            return Response({'msg': 'ok', 'data': res})
        except Exception as e:
            # 捕获并处理异常
            print(f"生成回复时发生错误: {e}")
            return Response({'msg': 'error', 'data': str(e)}, status=500)
